package com.zhang.antlr4demo.antlrtutorial;

import org.antlr.v4.runtime.CharStream;
import org.antlr.v4.runtime.CharStreams;
import org.antlr.v4.runtime.CodePointCharStream;
import org.antlr.v4.runtime.CommonTokenStream;

import java.util.HashMap;
import java.util.Map;

public class Calculator extends CalculatorBaseVisitor<Integer> {

    Map<String, Integer> memory = new HashMap<>();

    @Override
    public Integer visitAssign(CalculatorParser.AssignContext ctx) {
        String id = ctx.ID().getText();
        Integer value = visit(ctx.expression());
        memory.put(id, value);
        return value;
    }

    @Override
    public Integer visitPrintExpression(CalculatorParser.PrintExpressionContext ctx) {
        Integer value = visit(ctx.expression());
        System.out.println(value);
        return value;
    }

    @Override
    public Integer visitInt(CalculatorParser.IntContext ctx) {
        return Integer.parseInt(ctx.INT().getText());
    }

    @Override
    public Integer visitParents(CalculatorParser.ParentsContext ctx) {
        return visit(ctx.expression());
    }

    @Override
    public Integer visitMulDiv(CalculatorParser.MulDivContext ctx) {
        Integer left = visit(ctx.expression(0));
        Integer right = visit(ctx.expression(1));
        if (ctx.op.getType() == CalculatorParser.MUL) {
            return left * right;
        } else {
            return left / right;
        }
    }

    @Override
    public Integer visitAddSub(CalculatorParser.AddSubContext ctx) {
        Integer left = visit(ctx.expression(0));
        Integer right = visit(ctx.expression(1));
        if (ctx.op.getType() == CalculatorParser.ADD) {
            return left + right;
        } else {
            return left - right;
        }
    }

    @Override
    public Integer visitId(CalculatorParser.IdContext ctx) {
        String id = ctx.ID().getText();
        if (memory.containsKey(id)) return memory.get(id);
        return 0;
    }

    public static void main(String[] args) {
        String s = "1+2\n" +
                "a=1\n" +
                "b=2\n" +
                "a+b\n";

        CodePointCharStream charStream = CharStreams.fromString(s);

        CalculatorLexer lexer = new CalculatorLexer(charStream);

        CommonTokenStream tokens = new CommonTokenStream(lexer);
        System.out.println(tokens);

        CalculatorParser parser = new CalculatorParser(tokens);
        CalculatorParser.ProgramContext tree = parser.program();

        Calculator calculator = new Calculator();
        calculator.visit(tree);



    }

}
